glikelihood <- function(X, mu0, k0, alpha, betaX){
  ##Marginal likelihood of the data (also called "evidence"): P(X|mu0, k0, alpha, beta) - it is integral of P(X|m, v)p(m|v)p(v) with respect to m and v.
  N <- length(X)
  beta1 <- betaX + 1/2*sum((X-mean(X))^2) + 1/2*k0*N*(mean(X) - mu0)^2/(k0+N)
  -N/2*log(2 * pi) + 1/2*log(k0/(k0+N)) + alpha*log(betaX)-lgamma(alpha) + lgamma(alpha+N/2)-(alpha+N/2)*log(beta1)
}

mlikelihood <- function(D, clow, chigh){
  ##wrapper function that finds the marginal likelihood for each model that we are considering.
  ##browser()
  X<-diff(sqrt(D[2,]^2+D[3,]^2))
  dist <- mean(sqrt(D[2,]^2+D[3,]^2))
  #dist <- sqrt(D[2,]^2+D[3,]^2)

  ##Play around with the guys underneath. It would be even better if you could correct the code such that they are calculated automatically and then passed to this function.

  ##for dist derivative
  k0<-25*25;  #alpha=1; beta=1e-3
  k1<-25
  ##weight the two possible models according to the average distance between the actors in this segment of time: we're going to parameterize two Gammas via their mean and standard deviation.
  ## m1<-0.5; sd1<-0.5 
  ## m2<-5; sd2<-2.5
  ## w1 <- log(dgamma(dist, m1^2/sd1^2, m1/sd1^2))
  ## w2 <- log(dgamma(dist, m2^2/sd2^2, m2/sd2^2))

  #w1 <- log(dexp(dist, 2)); w2 <- log(dexp(dist, 1/5))
  w1 <- dexp(dist, 1); w2 <- dexp(dist, 1/10); s<-w1+w2; w1 <- log(w1/s); w2 <- log(w2/s)
  #w1 <- sum(log(dexp(dist, 1/1.9))); w2 <- sum(log(dexp(dist, 1/2.1)))  
  #w1=1; w2=1
  
  c(w1+glikelihood(X=X, mu0=0, k0=k0, alpha=alpha, betaX=betaX), w2+glikelihood(X=X, mu0=0, k0=k1, alpha=alpha, betaX=betaX))
}

toprob <- function(lprobs){
  ##converts unnormalised log probabilities to normalised true probabilities.
  probs <- exp(lprobs - max(lprobs))
  probs/sum(probs)
}

birth <- function(D, x, lambda, poisstrunc, mlikelihood, T){
  ##proposes a changepoint randomly in [0,T]. Then chooses two models, for the left and right of the new change-point, keeping all the other changepoints and models between changepoints as they are.
  cnew <- runif(1, min=0, max=T)
  mid <- max(which(c(0, x$cps ) <= cnew))
  clow <- c(0, x$cps)[mid]; chigh <- c(0, x$cps, T)[mid+1] 
  Dint <- as.matrix(D[, (D[1,]>clow) & (D[1,]<=chigh)]);
  Dlow <- as.matrix(Dint[, (Dint[1,]>clow) & (Dint[1,]<=cnew)])
  Dhigh <- as.matrix(Dint[, (Dint[1,]>cnew) & (Dint[1,]<=chigh)])
  if ((dim(Dlow)[2]< 2) | (dim(Dhigh)[2]< 2)){list(prob=2, x=x)}
  else {
    eint <- mlikelihood(Dint, clow, chigh)  
    elow <- mlikelihood(Dlow, clow, cnew);
    mlow <- sample(length(elow), prob=toprob(elow), size=1)
    ehigh <- mlikelihood(Dhigh, cnew, chigh);
    ##if (i == 110){browser()}
    mhigh <- sample(length(ehigh), prob=toprob(ehigh), size=1)
    cps <- sort(c(cnew, x$cps))
    ms <- c(x$ms, mlow, mhigh); is<- c(3*(1:length(x$ms)), 3*mid+1, 3*mid+2); ms <- ms[order(is)]; ms <- ms[-mid]
    newx <- list(cps=cps, ms=ms)
    ##underneath is the most important bit. We are going to accept the proposed birth (i.e. the new changepoint and the two new models on either side) with probability P(D|new x)P(new x) / [P(D| old x) P(old x)] * P(proposing old x from new x)/P(proposing new x from old x)*Jacobian. I am using equation 7 in Green's paper - it turns out that the Jacobian in there (the bit between the ||) evaluates to one in this example. (Everything is calculated on the log-scale for stability.)
    prob <- exp(log(lambda)-log(length(eint))   +elow[mlow]+ehigh[mhigh]-eint[x$ms[mid]]    -log(length(cps))+log(toprob(eint)[x$ms[mid]])   +log(T)-log(toprob(elow)[mlow])-log(toprob(ehigh)[mhigh])    +log(prob.move(newx, poisstrunc)["death"]-log(prob.move(x, poisstrunc)["birth"])))
    list(prob=prob, x=newx)
  }
}

death <- function(D, x, lambda, poisstrunc, mlikelihood, T){  
  ##randomly selects a changepoint and deletes it. Then picks a model for the aggregated region.
  cpid <- sample(1:length(x$cps), size=1)
  cp <- x$cps[cpid]
  clow <- c(0, x$cps)[cpid]
  chigh <- c(0, x$cps, T)[cpid+2]

  Dint <- as.matrix(D[, (D[1,]>clow) & (D[1,]<=chigh)]);
  Dlow <- as.matrix(Dint[, (Dint[1,]>clow) & (Dint[1,]<=cp)]);
  Dhigh <- as.matrix(Dint[, (Dint[1,]>cp) & (Dint[1,]<=chigh)]);

  eint <- mlikelihood(Dint, clow, chigh)
  mint <- sample(length(eint), prob=toprob(eint), size=1)
  elow <- mlikelihood(Dlow, clow, cp);
  ehigh <- mlikelihood(Dhigh, cp, chigh);

  cps <- x$cps[-cpid]
  ms <- x$ms; ms[cpid] <- mint; ms <- ms[-(cpid+1)]
  newx <- list(cps=cps, ms=ms)

  ##As with the birth we figure out what the acceptance probability should be. As above it is the prior ratio times the likelihood ratio times the probability of "coming back" divided by the probability of "going there", and finally times the Jacobian. Once again the Jacobian turns out to be 1. 
  prob <- exp(-log(lambda)+log(length(eint))   +eint[mint]-elow[x$ms[cpid]]-ehigh[x$ms[cpid+1]]    -log(T)+log(toprob(elow)[x$ms[cpid]])+log(toprob(ehigh)[x$ms[cpid+1]])     +length(x$cps)-log(toprob(eint)[mint]) + log(prob.move(newx, poisstrunc)["birth"])-log(prob.move(x, poisstrunc)["death"]))
  list(prob=prob, x=newx)
}

shift <- function(D, x, lamda, poisstrunc, mlikelihood, T){
  ##Pick a changepoint randomly, move it a bit, and repick the model on either side of it.
  cpid <- sample(1:length(x$cps), size=1)
  cp <- x$cps[cpid]
  clow <- c(0, x$cps)[cpid]
  chigh <- c(0, x$cps, T)[cpid+2]

  Dint <- as.matrix(D[, (D[1,]>clow) & (D[1,]<=chigh)]);
  
  cnew <- runif(1, min=clow, max=chigh)
  Dlownew <- as.matrix(Dint[, (Dint[1,]>clow) & (Dint[1,]<=cnew)])
  Dhighnew <- as.matrix(Dint[, (Dint[1,]>cnew) & (Dint[1,]<=chigh)])
  if ((dim(Dlownew)[2] <2) | (dim(Dhighnew)[2]<2)){list(prob=1, x=x)}
  else {
    Dlow <- as.matrix(Dint[, (Dint[1,]>clow) & (Dint[1,]<=cp)]);
    Dhigh <- as.matrix(Dint[, (Dint[1,]>cp) & (Dint[1,]<=chigh)]);
    elow <- mlikelihood(Dlow, clow, cp);
    ehigh <- mlikelihood(Dhigh, cp, chigh);

    elownew <- mlikelihood(Dlownew, clow, cnew);
    ehighnew <- mlikelihood(Dhighnew, cnew, chigh);

    newm <- c(sample(length(elownew), prob=toprob(elownew), size=1), sample(length(ehighnew), prob=toprob(ehighnew), size=1))
    cps<- x$cps; cps[cpid] <- cnew    
    ms<-x$ms; ms[c(cpid, cpid+1)] <- newm
    newx <- list(cps=cps, ms=ms)
    ##With birth and death, the new x that is proposed has a different dimension from the old x (i.e. basically their lengths are different). Here there is no change of dimension, so no need for a Jacobian term (this had turned out to be 1 for both birth and death anyway but that's just lucky). Acceptance probability is P(D|new x)P(new x) / [P(D| old x) P(old x)] * P(proposing old x from new x)/P(proposing new x from old x).
    prob <- exp(elownew[newm[1]]+ehighnew[newm[2]] - elow[x$ms[cpid]]-ehigh[x$ms[cpid+1]] +log(toprob(elow)[x$ms[cpid]])+log(toprob(ehigh)[x$ms[cpid+1]])-log(toprob(elownew))[newm[1]]-log(toprob(ehighnew))[newm[2]])
    list(prob=prob, x=newx)
  }
}

toggle <- function(D, x, lamda, poisstrunc, mlikelihood, T){
  ##pick a regime, and propose to change it (here there's no need for any acceptance probability, the probability of each regime is known exactly, conditional on the changepoints)
  mid <- sample(1:length(x$ms), size=1)
  clow <- c(0, x$cps)[mid]
  chigh <- c(0, x$cps, T)[mid+1]

  Dint <- as.matrix(D[, (D[1,]>clow) & (D[1,]<=chigh)])
  eint <- mlikelihood(Dint, clow, chigh)
  newx <- x; newx$ms[mid] <- sample(length(eint), prob=toprob(eint), size=1)
  list(prob=2, x=newx)  
}


prob.move <- function(x, poisstrunc){
  ##probabilities of making each move depending on the current set of change-points and models between changepoints (stored in x)
  ##you can't kill or move a changepoint if there aren't any, and we don't want more than poisstrunc changepoints, so:
  if (length(x$cps) <= 0){c(birth=1, death=0, shift=0, toggle=1)}
  else if (length(x$cps) >= poisstrunc){c(birth=0, death=1, shift=1, toggle=1)}
  else {c(birth=1, death=1, shift=1, toggle=1)}
  ##don't worry that these aren't normalised probabilities (they get normalised automatically in all the places they're used). The line above, for example, is supposed to mean that all moves have equal probability, 1/4.
}

getsample <- function(D, N=1000, lambda=1){
  ##assumes that time goes from [0, T] and expects roughly lambda changepoints every unit of time. Poisstrunc (underneath) is an upper bound on the number of changepoints (this isn't obligatory, but if we don't expect more than a certain number of changes, we might as well say so)
  start.time <- D[1,1]
  D[1,] <- D[1,]-start.time; T <- D[1,length(D[1,])]; poisstrunc <- ceiling(2*lambda*T)
  ##set priors on variance empirically
  X<-diff(sqrt(D[2,]^2+D[3,]^2))
  m.emp<-mean(replicate(100, 1/(sd(sample(X, size=min(10, ceiling(length(X)/2)))))^2))
  sd.emp<-sd(replicate(100, 1/(sd(sample(X, size=min(10, ceiling(length(X)/2)))))^2))
  alpha<<-m.emp^2/sd.emp^2; betaX <<- m.emp/sd.emp^2
  #alpha<-m.emp^2/sd.emp^2; betaX <- m.emp/sd.emp^2
  props <- list(); props <-list(birth, death, shift, toggle)
  x <- list(cps = c(), ms = c(1))
  s <- list(); s[[1]] <- x
  i <<- 2
  while (i <= N){
    cat(i, ": ", x$cps, "\n", sep="")  
    ##last set of changepoints, and models between each regime
    x <- s[[length(s)]]
    ##randomly pick a move from x to a new x. Possibilities are birth of a new changepoint, death of an existing changepoint, shift of an existing changepoint, or change of regime (the respective functions that do this are listed in "props" above)
    prop <- (props[[sample(length(props), size=1, prob=prob.move(x, poisstrunc))]])(D, x, lambda, poisstrunc, mlikelihood, T)
    ##accept the new x with probability min(1, x$prob)
    if (runif(1) < prop$prob){s[[length(s)+1]] <- prop$x}
    ##otherwise keep the old x
    else {s[[length(s)+1]] <- x}
    i <<- i+1 
  }
##postprocessing
  for (i in 1:length(s)){if (length(s[[i]])>0){s[[i]]$cps=s[[i]]$cps+start.time}}
  s
}

##"Sum" all of the changepoint sequences

##Patrick's version: too slow
## psum <- function(s){
##   cps<-c(); ms<- c()
##   for (i in 1:length(s)){
##     cps <- c(cps, c(0,s[[i]]$cps)); ms <- c(ms, s[[i]]$ms)
##   }
##   cps.sorted <- sort(cps, index.return=TRUE)
##   cps.s <- cps.sorted$x; ix <- cps.sorted$ix
##   pr <- rep(0, length(ms))
##   for (i in (1:(length(pr)-1))){
##     end <- which(ix==(ix[i]+1))
##     if (length(end)==0 | cps[ix[i]+1]==0){pr[i:length(pr)] <- pr[i:length(pr)] + as.numeric(ms[ix[i]]==1)}
##     else {
##       pr[i:(end-1)] <- pr[i:(end-1)]+ as.numeric(ms[ix[i]]==1)
##     }
##   }
##   pr[length(pr)] <- pr[length(pr)]+as.numeric(ms[length(pr)]==1)

##   ##delete replicate change-points
##   i <- 1
##   while (i < length(cps.s)){
##     if (cps.s[i]==cps.s[i+1]) {cps.s<-cps.s[-i]; pr<-pr[-i]}
##     else {i=i+1}
##     if (length(cps.s) != length(pr)){browser()}
##   }
##   list(cps=cps.s, pr=pr)
## }

##Mark's version - improvement?
psum <- function(s){
  print(Sys.time())
  cps<-c(); ms<- c()
  for (i in 1:length(s)){
    cps <- c(cps, c(0,s[[i]]$cps)); ms <- c(ms, s[[i]]$ms)
  }
  #print(Sys.time())
  cps.sorted <- sort(cps, index.return=TRUE)
  cps.s <- cps.sorted$x; ix <- cps.sorted$ix
  cps2 <- sort(ix,index.return=TRUE)
  cps2.sorted <- cbind(cps2$x,cps2$ix)
  pr2 <- rep(0, length(ms))
  checklen = length(cps2.sorted[,1])
  print(Sys.time())
  #modify this: use start+end markers during loop body, then sum over the whole thing once (instead of repeatedly adding to pieces)
  #sum+=(start-end) cell+=sum
  prstarts <- rep(0,length(pr2))
  prends <- rep(0,length(pr2))
  for (i in (1:(length(pr2)-1))){
    if(ix[i]+1 > checklen | cps[ix[i]+1] == 0) {
      prstarts[i] <- prstarts[i] + as.numeric(ms[ix[i]]==1)
      #pr2[i:length(pr2)] <- pr2[i:length(pr2)] + as.numeric(ms[ix[i]]==1)
    }
    else {
      end2 <- cps2.sorted[ix[i]+1,2]
      prstarts[i] <- prstarts[i] + as.numeric(ms[ix[i]]==1)
      prends[end2] <- prends[end2] + as.numeric(ms[ix[i]]==1)
      #pr2[i:(end2-1)] <- pr2[i:(end2-1)]+ as.numeric(ms[ix[i]]==1)
    }
  }
  prsums <- rep(0,length(prstarts))
  csum <- 0
  prstarts[length(prstarts)] <- prstarts[length(prstarts)]+as.numeric(ms[length(prstarts)]==1)
  #print(length(prstarts))
  for(i in 1:length(prstarts)) {
    csum <- csum + prstarts[i]
    csum <- csum - prends[i]
    #if(i == length(prstarts)) {
    #  print(prstarts[i])
    #  print(prends[i])
    #}
    prsums[i] <- csum
  }
  pr2 <- prsums
  print(Sys.time())
  #pr2[length(pr2)] <- pr2[length(pr2)]+as.numeric(ms[length(pr2)]==1)
  #print(prsums)
  #print(pr2)
  #for(i in 1:length(prstarts)) {
  #  if(prsums[i] != pr2[i]) {
  #    print("OH NO!")
  #    print(i)
  #    print(prsums[i])
  #    print(pr2[i])
  #  }
  #}
  #cps25 <- cps.s
  cps3 <- c()
  #pr25 <- pr2
  pr3 <- c()
  i <- 1
  while(i <= length(cps.s)) {
    if(i == length(cps.s) | cps.s[i] != cps.s[i+1]) {
      cps3[length(cps3)+1] <- cps.s[i]
      pr3[length(pr3)+1] <- pr2[i]
    }
    i = i+1
  }
  #print(pr3)
  print(Sys.time())
  list(cps=cps3, pr=pr3)
}
